﻿#include "GeneratedMesh.h"
#include <vtkm/cont/DataSetBuilderUniform.h>
#include <vtkm/worklet/WorkletMapTopology.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/filter/Triangulate.h>

RegisterObject(GeneratedMesh);

struct ComputeTopolgy : vtkm::worklet::WorkletVisitCellsWithPoints
{
  using ControlSignature = void(CellSetIn, FieldInPoint, FieldOut, FieldOut);
  using ExecutionSignature = void(_2, _3, _4);

  template<typename PType>
  VTKM_EXEC void operator()(const PType& pts, vtkm::Vec3f& centroid, vtkm::Vec3f& spacing) const
  {
    if (pts.GetNumberOfComponents() == 4)
    {
      centroid = (pts[0] + pts[1] + pts[2] + pts[3]) / 4;
      spacing[0] = (pts[1] - pts[0])[0];
      spacing[1] = (pts[2] - pts[1])[1];
    }
    else if (pts.GetNumberOfComponents() == 8)
    {
      centroid = (pts[0] + pts[1] + pts[2] + pts[3] + pts[4] + pts[5] + pts[6] + pts[7]) / 8;
      spacing[0] = (pts[1] - pts[0])[0];
      spacing[1] = (pts[2] - pts[1])[1];
      spacing[2] = (pts[4] - pts[0])[2];
    }
    else
    {
      RaiseError("unfinished");
    }
  }
};


struct BuildBndList : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn point, FieldOut list);
  using ExecutionSignature = void(_1, _2);

  template<typename PType>
  VTKM_EXEC void operator()(const PType& p, vtkm::Id& bid) const
  {
    if (vtkm::Abs(p[0] - _x_min) < _tol) // left
    {
      bid = 3;
    }
    else if (vtkm::Abs(p[0] - _x_max) < _tol) // right
    {
      bid = 1;
    }
    else if (vtkm::Abs(p[1] - _y_min) < _tol) // bottom
    {
      bid = 0;
    }
    else if (vtkm::Abs(p[1] - _y_max) < _tol) // top
    {
      bid = 2;
    }
    else // 内部点
    {
      bid = -1;
    }

    if (vtkm::Abs(p[0] - _x_min) < _tol && vtkm::Abs(p[1] - _y_min) < _tol)
    {
      bid = 4;
    }
    else if (vtkm::Abs(p[0] - _x_min) < _tol && vtkm::Abs(p[1] - _y_max) < _tol)
    {
      bid = 4;
    }
    else if (vtkm::Abs(p[0] - _x_max) < _tol && vtkm::Abs(p[1] - _y_min) < _tol)
    {
      bid = 4;
    }
    else if (vtkm::Abs(p[0] - _x_max) < _tol && vtkm::Abs(p[1] - _y_max) < _tol)
    {
      bid = 4;
    }
  }

private:
  vtkm::FloatDefault _x_min = 0, _x_max = 1, _y_min = 0, _y_max = 1;
  vtkm::FloatDefault _tol = 1E-06;
};

struct CoodConvert : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn in, FieldOut out);
  using ExecutionSignature = void(_1, _2);
  template<typename InType, typename OutType>
  VTKM_EXEC void operator()(const InType& in, OutType& out) const
  {
    out = in;
  }
};

GeneratedMesh::GeneratedMesh(const Configuration& cfg)
  : Mesh(cfg)
  , _triangularization(Get<bool>("triangularization", false))
{
  _dims = GetVectorValue<int>("dims");

  if (_dims.size() == 1)
  {
    Error("一维网格生成不支持");
  }
  else if (_dims.size() == 2)
  {
    _x_ranges = GetVectorValue<Real>("x_range");
    _y_ranges = GetVectorValue<Real>("y_range");
    _nx = _dims[0];
    _ny = _dims[1];
  }
  else if (_dims.size() == 3)
  {
    _x_ranges = GetVectorValue<Real>("x_range");
    _y_ranges = GetVectorValue<Real>("y_range");
    _z_ranges = GetVectorValue<Real>("z_range");
  }
  else
  {
    Error("dim设置错误");
  }
}

void GeneratedMesh::Build()
{
  Vec2f start{ _x_ranges[0], _y_ranges[0] };
  Vec2f end{ _x_ranges[1], _y_ranges[1] };
  Vec2f space = end - start;
  Id2 point_dim = { _dims[0] + 1, _dims[1] + 1 };

  for (size_t i = 0; i < 2; i++)
    space[i] = (end[i] - start[i]) / _dims[i];

  vtkm::cont::DataSet data_set = vtkm::cont::DataSetBuilderUniform::Create(point_dim, start, space);
  this->_cell_set = data_set.GetCellSet();
  this->_coord = data_set.GetCoordinateSystem();

  vtkm::cont::Invoker{}(ComputeTopolgy{}, _cell_set, _coord, _centroid, _spacing);
  vtkm::cont::Invoker{}(BuildBndList{}, _coord, _boundary_id);

  if (_triangularization)
  {
    vtkm::filter::Triangulate triangulateFilter;
    vtkm::cont::DataSet output = triangulateFilter.Execute(data_set);
    this->_cell_set = output.GetCellSet();
    ArrayHandle<Vec3f> pts;
    Invoker{}(CoodConvert{}, output.GetCoordinateSystem(), pts);
    this->_coord = CoordinateSystem("coord", pts);
  }

}
